import torch
import sys
sys.path.append('..')
from Function.readdatafunction import readdata
from Function.CNN import SimpleCNN
from Function.loss import train, cross_entropy_normal

torch.cuda.empty_cache()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Xtrain, Ytrain, X, Y, res = readdata("../dataset/Salinas/Salinas_corrected.mat", "salinas_corrected",
                                        "../dataset/Salinas/Salinas_gt.mat", "salinas_gt",
                                        17, 204, 16, "CNN")
Xtrain = torch.from_numpy(Xtrain)
Ytrain = torch.from_numpy(Ytrain).long()

net = SimpleCNN(num_classes=16, in_channels=204)
PATH = './model/Salinas/CNN_Salinas.pt'
net.load_state_dict(torch.load(PATH))
net.to(device)

lr, num_epochs = 0.00001, 100
batch_size = 64

train(net, Xtrain, Ytrain, num_epochs, lr, batch_size, cross_entropy_normal, device)
print("OK")
torch.save(net.state_dict(), PATH)  # 推荐的文件后缀名是pt或pth